#include "trace.h"

#include <float.h>

#include "../cuCommon/cuError.h"

//////////////////////////////////////////////////////////////////////////
/// Kernels definitions
//////////////////////////////////////////////////////////////////////////

#define CU_START_KERNEL_2D(w)															 \
	int tidx = threadIdx.x + blockDim.x * blockIdx.x;				 \
  int tidy = threadIdx.y + blockDim.y * blockIdx.y;				 \
	int tid = tidx + tidy * (w);

__global__
void cuGenerateEyeRays_Perspective_kernel(IN  int2     resolution, // real pixel plane
																					IN  float2   scale,      // pixel scale
																					IN  int2     sample,     // aa mode
																					IN  float3   eye,        // camera eye
																					IN  float3   dir,        // camera direction
																					IN  float3   right,      // camera right
																					IN  float3   up,         // camera up
																					IN  float    zn,         // camera z near distance
																					OUT float4 * pRayO,      // ray origin buffer
																					OUT float4 * pRayD,      // ray direction buffer
																					IN  uint32   nWidth)     // width of a threads grid
{
	CU_START_KERNEL_2D(nWidth);

	if (tidx < resolution.x && tidy < resolution.y)
	{
		for (float sx = 0; sx < sample.x; sx++)
			for (float sy = 0; sy < sample.y; sy++)
			{
				float2 sp = { (sx + 0.5f) / sample.x , (sy + 0.5f) / sample.y };

				// pixel plane
				float2 pp;
				pp.x = scale.x * (tidx - 0.5f * resolution.x + sp.x);
				pp.y = scale.y * (tidy - 0.5f * resolution.y + sp.y);

				int offset = sx + sy * sample.x;

				pRayO[tid + offset * nWidth] = f4(eye);
				pRayD[tid + offset * nWidth] = f4(normalize(right * pp.x + up * pp.y + dir * zn));
			}
	}
}


__global__
void cuGenerateEyeRays_Orthogonal_kernel(int2     resolution,  // resolution
																         float2   scale,       // scale plane
																				 int2     sample,      // aa sample mode
																				 float3   eye,				 // camera eye
																				 float3   dir,         // camera direction
																				 float3   right,       // camera right
																				 float3   up,          // camera up
																				 float    zn,          // camera z near distance
																				 float4 * pRayO,       // ray origin buffer
																				 float4 * pRayD,       // ray direction buffer
																				 uint32   nWidth)      // width of a threads grid
{
	CU_START_KERNEL_2D(nWidth);

	if (tidx < resolution.x && tidy < resolution.y)
	{
		for (float sx = 0; sx < sample.x; sx++)
				for (float sy = 0; sy < sample.y; sy++)
				{
					float2 sp = { (sx + 0.5f) / sample.x , (sy + 0.5f) / sample.y };

					//sp = rnUnitSquare();

					float2 pp;
					pp.x = scale.x * (tidx - 0.5f * resolution.x + sp.x);
					pp.y = scale.y * (tidy - 0.5f * resolution.y + sp.y);
					
					int offset = sx + sy * sample.x;

					pRayO[tid + offset * nWidth] = f4(eye + right * pp.x + up * pp.y);
					pRayD[tid + offset * nWidth] = f4(dir);
				}
	}
}

#define BLOCK_SIZE 64
#define V0_x s_vert_pos[threadIdx.x+ 0*BLOCK_SIZE]
#define V0_y s_vert_pos[threadIdx.x+ 1*BLOCK_SIZE]
#define V0_z s_vert_pos[threadIdx.x+ 2*BLOCK_SIZE]

#define V1_x s_vert_pos[threadIdx.x+ 3*BLOCK_SIZE]
#define V1_y s_vert_pos[threadIdx.x+ 4*BLOCK_SIZE]
#define V1_z s_vert_pos[threadIdx.x+ 5*BLOCK_SIZE]

#define V2_x s_vert_pos[threadIdx.x+ 6*BLOCK_SIZE]
#define V2_y s_vert_pos[threadIdx.x+ 7*BLOCK_SIZE]
#define V2_z s_vert_pos[threadIdx.x+ 8*BLOCK_SIZE]

#define E1_x s_vert_pos[threadIdx.x+ 9*BLOCK_SIZE]
#define E1_y s_vert_pos[threadIdx.x+10*BLOCK_SIZE]
#define E1_z s_vert_pos[threadIdx.x+11*BLOCK_SIZE]

#define E2_x s_vert_pos[threadIdx.x+12*BLOCK_SIZE]
#define E2_y s_vert_pos[threadIdx.x+13*BLOCK_SIZE]
#define E2_z s_vert_pos[threadIdx.x+14*BLOCK_SIZE]

//////////////////////////////////////////////////////////////////////////
/// texture references for trace kernel
//////////////////////////////////////////////////////////////////////////
texture<float4, 1, cudaReadModeElementType> g_texPosition;
texture<int1,   1, cudaReadModeElementType> g_texPositionIdx; 

__global__
void cuTrace_kernel(IN  float4 * pRayO,
										IN  float4 * pRayD, 
										IN  uint32   nIdx,
										IN  uint32   nPosO,
										IN  uint32   nPosIdxO,
										OUT float  * pTmin,
										OUT int2   * pID,
										OUT uchar4 * pRGBA8,
										IN  uint32   iMesh,
										IN  int32  * pIdx,
										IN  float4 * pPos,
										IN  uint32   nRay)
{
	CU_START_KERNEL_2D(nRay);

	//__shared__ float s_vert_pos[15 * BLOCK_SIZE];

	if (tid < nRay)
	{
		SRay ray;
		float4 rayO = pRayO[tid];
		float4 rayD = pRayD[tid];

		ray.O = xyz(rayO);
		ray.D = xyz(rayD);

		//float3 tbg = { FLT_MAX, -1, -1 };
		float tvalue = FLT_MAX;

		float  tminprev = pTmin[tid];
		uchar4 color    = pRGBA8[tid];
		int2   id       = pID[tid];

		for (int iIdx = 0; iIdx < nIdx; iIdx+=3)
		{
			int3 i3;
			i3.x = tex1Dfetch(g_texPositionIdx, nPosIdxO + iIdx+0).x;
			i3.y = tex1Dfetch(g_texPositionIdx, nPosIdxO + iIdx+1).x;
			i3.z = tex1Dfetch(g_texPositionIdx, nPosIdxO + iIdx+2).x;

			STriangle tri;
			tri.m_V[0] = xyz(tex1Dfetch(g_texPosition, nPosO + i3.x));
			tri.m_V[1] = xyz(tex1Dfetch(g_texPosition, nPosO + i3.y));
			tri.m_V[2] = xyz(tex1Dfetch(g_texPosition, nPosO + i3.z));

			if (Intersect(tri, ray, tvalue))
			{
				if (tvalue < tminprev)
				{
					tminprev = tvalue;
					id.x = iMesh;
					id.y = iIdx;

					color = uc4(255);
				}
			}
		}

		pRGBA8[tid] = color;
		pTmin[tid]  = tminprev;
		pID[tid]    = id;
	}
}

__shared__ extern float smem_TraceKdTree_sf_[];
__shared__ extern uchar smem_TraceKdTree_uc_[];

#define KD_NODE_COUNT(n)          (as_uint(n.x) >> 2)
#define KD_NODE_SPLIT(n)          (as_float(as_uint(n.x) & (~TREE_LEAF)))
#define KD_NODE_AXIS(n)           (as_uint(n.x) & TREE_LEAF)
#define KD_NODE_NEXT(n)           (n.y)
#define KD_NODE_START(n)          (n.y)
#define KD_NODE_T(n)              (n.x)

#define KD_TRACE_RAY_O_AXIS(i)    (smem_TraceKdTree_sf_[(i+0*1) * blockDim.x * (blockDim.y + threadIdx.y) + threadIdx.x])
#define KD_TRACE_INV_D_AXIS(i)    (smem_TraceKdTree_sf_[(i+3*1) * blockDim.x * (blockDim.y + threadIdx.y) + threadIdx.x])
#define KD_TRACE_SGN_D_AXIS(i)    (smem_TraceKdTree_uc_[(i+6*4) * blockDim.x * (blockDim.y + threadIdx.y) + threadIdx.x])

texture<float2, 1, cudaReadModeElementType> g_Nodes;
texture<float4, 1, cudaReadModeElementType> g_Position;
texture<float2, 1, cudaReadModeElementType> g_Face;
texture<float2, 1, cudaReadModeElementType> g_Index;

__global__
void cuTraceKdTree_kernel(IN  float2       * pNodes,
													IN  S_Box         box,
													IN  float4       * pRayO,
													IN  float4       * pRayD,
													IN  uint32         nRay,
													IN  float4       * pPosition,
													IN  int32        * pFace,
													IN  uint32       * pIndex,
													OUT float        * pTmin, 
													OUT int2         * pId, 
													OUT uchar4       * pRgba8)
{
	CU_START_KERNEL_2D(nRay);

	int iRay = tid;

	float2 node_stack[32];

	float tnear = 0.0f, tfar = 0.0f, split = 0.0f, tsplit = 0.0f, thit = FLT_MAX, tprev = thit;
	int stack_index = 0;
	int node_index = 0;
	int axis = TREE_LEAF;

	float3 rayD = xyz(pRayD[iRay]);
	float3 rayO = xyz(pRayO[iRay]);

	int2 id = { -1, -1 };

	bool bHit = true;

	KD_TRACE_RAY_O_AXIS(0) = rayO.x;
	KD_TRACE_RAY_O_AXIS(1) = rayO.y;
	KD_TRACE_RAY_O_AXIS(2) = rayO.z;

	KD_TRACE_INV_D_AXIS(0) = 1.0f / rayD.x;
	KD_TRACE_INV_D_AXIS(1) = 1.0f / rayD.y;
	KD_TRACE_INV_D_AXIS(2) = 1.0f / rayD.z;

	KD_TRACE_SGN_D_AXIS(0) = (rayD.x >= 0) ? 1 : 0;
	KD_TRACE_SGN_D_AXIS(1) = (rayD.y >= 0) ? 1 : 0;
	KD_TRACE_SGN_D_AXIS(2) = (rayD.z >= 0) ? 1 : 0;

	if (Intersect(box.pmin, box.pmax, rayO, rayD, tnear, tfar))
	{
		if (tnear < 0) tnear = 0.0f;
		bHit = false;
	}

	while( !bHit )
	{
		split = KD_NODE_SPLIT( pNodes[node_index] );
		axis  = KD_NODE_AXIS( pNodes[node_index] );

		tsplit = (split - KD_TRACE_RAY_O_AXIS(axis)) * KD_TRACE_INV_D_AXIS(axis);

		int first_left = KD_TRACE_SGN_D_AXIS(axis);

		int near = 0;
		int far = 0;

		int left = node_index+1;
		int right = KD_NODE_NEXT( pNodes[node_index] );

		if (axis != TREE_LEAF)
		{
			if (first_left)
			{
				near = left;
				far = right;
			}
			else
			{
				far = left;
				near = right;
			}

			if(tsplit < tnear)
			{
				node_index = far;
			}
			else if (tsplit > tfar)
			{
				node_index = near;
			}
			else
			{
				node_index = near;
				KD_NODE_T( node_stack[stack_index] ) = tfar;
				KD_NODE_NEXT( node_stack[stack_index] ) = far;

				stack_index++;

				tfar = tsplit;
			}
		}
		else
		{
			int32 prim_start = KD_NODE_START( pNodes[node_index] );
			int32 prim_count = KD_NODE_COUNT( pNodes[node_index] );

			for (int32 iPrimitive = 0; iPrimitive < prim_count; iPrimitive++)
			{
				int32 face_idx = pIndex[ prim_start + iPrimitive ];

				int3 face = i3( pFace[3*face_idx], pFace[3*face_idx+1], pFace[3*face_idx+2] );

				float3 v[3] = { xyz(pPosition[face.x]), xyz(pPosition[face.y]), xyz(pPosition[face.z]) };

				if (Intersect(v[0], v[1], v[2], rayO, rayD, thit))
				{					
					if (tprev > thit)
					{
						const uchar4 white = {255, 255, 255, 255};

						pRgba8[iRay] = white;
						tprev = thit;					
						
						pId[iRay] = i2(face_idx, 0);

						bHit = true;
					} // (tprev > thit)
				}  // if (Intersect(v[0], v[1], v[2], rayO, rayD, thit))
			} // for (int32 iPrimitive = 0; iPrimitive < prim_count; iPrimitive++)

			if (!bHit && stack_index > 0) 
			{
				stack_index--;

				tnear = tfar;
				tfar = KD_NODE_T( node_stack[stack_index] );
				node_index = KD_NODE_NEXT( node_stack[stack_index] );
			}
			else
			{
				pTmin[iRay] = tprev;
				bHit = true;
			}
		} //if (axis != KD_LEAF)
	} // while (!bhit)
}

__global__
void cuTraceKdTree_kernel_tex(IN  float2       * pNodes,
															IN  S_Box         box,
															IN  float4       * pRayO,
															IN  float4       * pRayD,
															IN  uint32         nRay,
															IN  float4       * pPosition,
															IN  int32        * pFace,
															IN  uint32       * pIndex,
															OUT float        * pTmin, 
															OUT int2         * pId, 
															OUT uchar4       * pRgba8)
{
	CU_START_KERNEL_2D(nRay);

	int iRay = tid;

	float2 node_stack[32];

	float tnear = 0.0f, tfar = 0.0f, split = 0.0f, tsplit = 0.0f, thit = FLT_MAX, tprev = thit;
	int stack_index = 0;
	int node_index = 0;
	int axis = TREE_LEAF;

	float3 rayD = xyz(pRayD[iRay]);
	float3 rayO = xyz(pRayO[iRay]);

	int2 id = { -1, -1 };

	bool bHit = true;

	KD_TRACE_RAY_O_AXIS(0) = rayO.x;
	KD_TRACE_RAY_O_AXIS(1) = rayO.y;
	KD_TRACE_RAY_O_AXIS(2) = rayO.z;

	KD_TRACE_INV_D_AXIS(0) = 1.0f / rayD.x;
	KD_TRACE_INV_D_AXIS(1) = 1.0f / rayD.y;
	KD_TRACE_INV_D_AXIS(2) = 1.0f / rayD.z;

	KD_TRACE_SGN_D_AXIS(0) = (rayD.x >= 0) ? 1 : 0;
	KD_TRACE_SGN_D_AXIS(1) = (rayD.y >= 0) ? 1 : 0;
	KD_TRACE_SGN_D_AXIS(2) = (rayD.z >= 0) ? 1 : 0;

	if (Intersect(box.pmin, box.pmax, rayO, rayD, tnear, tfar))
	{
		if (tnear < 0) tnear = 0.0f;
		bHit = false;
	}

	while( !bHit )
	{
		float2 curr_node = tex1Dfetch(g_Nodes, node_index);
		
		split = KD_NODE_SPLIT( curr_node );
		axis  = KD_NODE_AXIS( curr_node );

		tsplit = (split - KD_TRACE_RAY_O_AXIS(axis)) * KD_TRACE_INV_D_AXIS(axis);

		int first_left = KD_TRACE_SGN_D_AXIS(axis);

		int near = 0;
		int far = 0;

		int left = node_index+1;
		int right = KD_NODE_NEXT( curr_node );

		if (axis != TREE_LEAF)
		{
			if (first_left)
			{
				near = left;
				far = right;
			}
			else
			{
				far = left;
				near = right;
			}

			if(tsplit < tnear)
			{
				node_index = far;
			}
			else if (tsplit > tfar)
			{
				node_index = near;
			}
			else
			{
				node_index = near;
				KD_NODE_T( node_stack[stack_index] ) = tfar;
				KD_NODE_NEXT( node_stack[stack_index] ) = far;

				stack_index++;

				tfar = tsplit;
			}
		}
		else
		{
			int32 prim_start = KD_NODE_START( curr_node );
			int32 prim_count = KD_NODE_COUNT( curr_node );

			for (int32 iPrimitive = 0; iPrimitive < prim_count; iPrimitive++)
			{
				int32 face_idx = pIndex[ prim_start + iPrimitive ];

				int3 face = i3( pFace[3*face_idx], pFace[3*face_idx+1], pFace[3*face_idx+2] );

				float3 v[3] = { xyz(pPosition[face.x]), xyz(pPosition[face.y]), xyz(pPosition[face.z]) };

				if (Intersect(v[0], v[1], v[2], rayO, rayD, thit))
				{					
					if (tprev > thit)
					{
						const uchar4 white = {255, 255, 255, 255};

						pRgba8[iRay] = white;
						tprev = thit;					
						
						pId[iRay] = i2(face_idx, 0);

						bHit = true;
					} // (tprev > thit)
				}  // if (Intersect(v[0], v[1], v[2], rayO, rayD, thit))
			} // for (int32 iPrimitive = 0; iPrimitive < prim_count; iPrimitive++)

			if (!bHit && stack_index > 0) 
			{
				stack_index--;

				tnear = tfar;
				tfar = KD_NODE_T( node_stack[stack_index] );
				node_index = KD_NODE_NEXT( node_stack[stack_index] );
			}
			else
			{
				pTmin[iRay] = tprev;
				bHit = true;
			}
		} //if (axis != KD_LEAF)
	} // while (!bhit)
}

//////////////////////////////////////////////////////////////////////////
/// Wrappers definitions
//////////////////////////////////////////////////////////////////////////

void cuGenerateEyeRays(IN  C_Camera  & camera, 
											 IN  int2     & resolution,
											 IN  int2     & sample,
											 IN  uint32     type,
											 OUT float4   * pRayO,
											 OUT float4   * pRayD)
{
	CBenchmark genraysBench;

	float3 eye = camera.Eye();
	float3 dir = camera.Dir();
	float3 up  = camera.Up();

	float  zn  = camera.zN();
	float  zf  = camera.zF();
	float  asp = camera.Asp();
	float  fov = camera.FoV();

	float3 right = camera.Right();

	float2 scale = { 1.0f, 1.0f };

	float2 vp; // virtual plane
	vp.x = zn * tanf(to_radian(fov * 0.5f));
	vp.y = vp.x / asp;

	if (type == RT_PERSPECTIVE)
	{
		genraysBench.Start("cuGenerateEyeRays Perspective");

		scale.x = scale.x / (resolution.x * 0.5f) * vp.x;
		scale.y = scale.y / (resolution.y * 0.5f) * vp.y;

		dim3 threads(32, 4);
		dim3 blocks((resolution.x + threads.x - 1) / threads.x, 
			          (resolution.y + threads.y - 1) / threads.y);

		cuGenerateEyeRays_Perspective_kernel<<<blocks, threads>>>(resolution, scale, sample, 
			                                                        eye, dir, right, up, zn,
																															pRayO, pRayD,
																															resolution.x);
		cudaThreadSynchronize();

		genraysBench.End();
		genraysBench.Print();
	}

	if (type == RT_ORTHOGONAL)
	{
		genraysBench.Start("mpGenerateEyeRays Orthogonal");

		scale.x = scale.x / (resolution.x * 0.5f) * vp.x;
		scale.y = scale.y / (resolution.y * 0.5f) * vp.y;

		dim3 threads(32, 4);
		dim3 blocks((resolution.x + threads.x - 1) / threads.x, 
			          (resolution.y + threads.y - 1) / threads.y);

		cuGenerateEyeRays_Orthogonal_kernel<<<blocks, threads>>>(resolution, scale, sample, 
																														 eye, dir, right, up, zn,
																														 pRayO, pRayD, 
																														 resolution.x);
		cudaThreadSynchronize();

		genraysBench.End();
		genraysBench.Print();
	}
}

void cuTrace(IN  float4  * pRayO,
						 IN  float4  * pRayD, 
						 IN  uint32    nRay,
						 IN  CglMesh * pMesh,
						 IN  uint32    nMesh,
						 OUT float   * pTmin, 
						 OUT int2    * pId, 
						 OUT uchar4  * pRgba8)
{
	CBenchmark traceBench;

	traceBench.Start("cuTrace");

	dim3 threads(64, 1);
	dim3 blocks( (nRay + threads.x - 1) / threads.x );
	blocks.y = nRay / (blocks.x * threads.x);

  cudaChannelFormatDesc posDesc    = cudaCreateChannelDesc<float4>();
	cudaChannelFormatDesc posIdxDesc = cudaCreateChannelDesc<int1>();

	for (uint32 iMesh = 0; iMesh < nMesh; iMesh++)
	{
		pMesh[iMesh].pos.Map();
		pMesh[iMesh].posIdx.Map();

		float4 * pPos    = pMesh[iMesh].pos.dPtrR()    + pMesh[iMesh].pos.offsetR();
		int    * pPosIdx = pMesh[iMesh].posIdx.dPtrR() + pMesh[iMesh].posIdx.offsetR();

		uint32 nPos    = pMesh[iMesh].pos.size();
		uint32 nPosIdx = pMesh[iMesh].posIdx.size();

		uint32 nPosO = 0;
		uint32 nPosIdxO = 0;

		cudaBindTexture((size_t *) &nPosO,    g_texPosition,    (void *) pPos,    nPos * sizeof(float4));
		cudaBindTexture((size_t *) &nPosIdxO, g_texPositionIdx, (void *) pPosIdx, nPosIdx * sizeof(int));

		//std::cout << "Position offset is \n" << nPosO << std::endl;
		//std::cout << "Position index offset is \n" << nPosIdxO << std::endl;

		nPosO = nPosO / sizeof(float4);
		nPosIdxO = nPosIdxO / sizeof(int);

		cuTrace_kernel<<<blocks, threads>>>(pRayO, pRayD, 
									                      nPosIdx, nPosO, nPosIdxO,
									                      pTmin, pId, pRgba8,
									                      iMesh, pPosIdx, pPos, nRay);

		cudaThreadSynchronize();

		pMesh[iMesh].pos.Unmap();
		pMesh[iMesh].posIdx.Unmap();
	}

	traceBench.End();
}

void cuTraceKdTree(IN  float2    * pKd,
									 IN  S_Box    & box, 
									 IN  float4    * pRayO,
									 IN  float4    * pRayD,
									 IN  uint32      nRay,
									 IN  float4    * pPosition,
									 IN  int32     * pFace,
									 IN  uint32    * pIndex,
									 OUT float     * pTmin, 
									 OUT int2      * pId, 
									 OUT uchar4    * pRgba8)
{
	CBenchmark traceBench;

	traceBench.Start("cuTraceKdTree");

	dim3 threads(64, 1);
	dim3 blocks( (nRay + threads.x - 1) / threads.x );
	blocks.y = nRay / (blocks.x * threads.x);

	cuTraceKdTree_kernel<<<blocks, threads>>>(pKd, box, pRayO, pRayD, nRay, pPosition, pFace, pIndex, pTmin, pId, pRgba8);

	cudaThreadSynchronize();

	traceBench.End();
}